%%%% The CIM function simulates a Coherent Ising Machine (CIM) 
%%%% based on the Degenerate Optical Parametric Oscillator (DOPO)
%%%% model using the quantum jump method. The simulation uses
%%%% the xSPDE 4.2 library to solve open quantum dynamics in a 
%%%% truncated Fock basis.

function [e,data,p] = CIM()
% Coherent Ising Machine (CIM) simulation using quantum jump method in xSPDE 4.2
% Returns expectation values `e`, raw data `data`, and parameter structure `p`

% --- Simulation Parameters ---
p.name       = 'CIM';                               % Name of the simulation
p.nm         = 2;                                   % Number of modes (spins)
p.N          = 15;                                  % Photon number cutoff
p.nmax       = (p.N+1)*ones(1,p.nm);                % Photon cutoff for each mode
p.g          = @(p) 0.6;                            % Nonlinear loss coefficient
p.alph       = sqrt(2.4)/0.6;                       % Coherent amplitude α = sqrt(λ)/g
p.steps      = 1;                                   % Internal time integration steps
p.lam        = @(p) p.g(p)^2*p.alph^2;              % Pump rate λ = g^2 * α^2
% --- Ising Coupling ---
J = coupling_matrix(p);                             % Generate Ising coupling matrix
p.jj         = @(p) 1;                              % Scaling factor for J
p.J          = @(p) p.jj(p)*J;                      % Coupling matrix function handle
% --- Time and Ensemble Configuration ---
p.ranges     = 2;                                   % Total simulation time (in scaled units)
p.points     = 400;                                 % Number of time points
p.ensembles  = [500,1,10];                          % [trajectories, observables, plots]
 % --- Simulation Method Flags ---
p.quantum    = 1;                                   % Use quantum simulation
p.jump       = 1;                                   % Enable quantum jump method
p.sparse     = 1;                                   % Use sparse matrices
p.method     = @RK4;                                % Time integration method
 % --- Operators and Dissipation ---
p.a          = mkbose(p);                           % Create bosonic annihilation operators
p.gamma{1}   = @(p) ones(1,p.nm);                   % Linear loss rates
p.gamma{2}   = @(p) p.g(p)^2*ones(1,p.nm);          % Nonlinear (two-photon) loss rates
p.gamma{3}   = @(p) abs(p.J(p));                    % Mode coupling loss rates
% Define jump operators (dissipators)
p.L          = {@(m,p) p.a{m}, @(m,p) p.a{m}^2};    % L1 = a, L2 = a^2
p.L{3}       = @(i,p) coupling_generation(i,p);     % L3 = a_i - J_ij a_j
% --- Hamiltonian and Initial State ---
p.H          = @Ham;                                % Hamiltonian function handle
p.initial    = @(~,p) mkvacs(p);                    % Initial state (vacuum or cat)
p.checks     = 1;                                   % Enable Time-Step checks
% --- Precompute Ising Problem Information ---
[p.E_Ising, p.total_spin_config] = Ising_config(p); % All spin configs and their energies
p.idx = 0;                                          % (Unused; placeholder)
[p.success_p, p.success_n] = success_prob(p.N);     % Projectors for up/down states
% --- Observables ---
% p.observe{1} = @(psi,p) purity(psi,p);            % (Optional) Purity observable
p.observe{2} = @(psi,p) success_rate(psi,p);        % Success probability observable
% --- Run the xSPDE Simulation ---
[e,data,p]  = xspde(p); 
end

%%%% DOPO Hamiltonian
%% --- DOPO Hamiltonian ---
function H = Ham(p)
    % Returns the effective Hamiltonian of the DOPO CIM system
    H=sparse((p.N+1)^p.nm,(p.N+1)^p.nm);
    for m=1:p.modes
       H = H+(p.a{m}')^2-p.a{m}^2;
    end
    H = 1i*p.lam(p)/2*H;
end

%% --- Success Rate Observable ---
function [suc_rate]=success_rate(psi,p)
    % Computes the overlap of the final wavefunction with minimum-energy spin configurations
    P_sc = 0;
    suc_rate = 0;
    min_energy = min(p.E_Ising); %% lowest energy 
    min_index = find(p.E_Ising == min_energy); %% spin config index that gives the lowest energy
    min_config = p.total_spin_config(min_index,:);
    for loop = 1:length(min_index)
        psi_update = psi;
        for j = 1:p.nm
            if j == 1
                if min_config(loop,j) == 1
                    S =  kron( p.success_p, speye((p.N+1)^(p.nm-1)) ) ;
                else
                    S =  kron( p.success_n, speye((p.N+1)^(p.nm-1)) ) ;
                end
            elseif j == p.nm
                if min_config(loop,j) == 1
                    S =  kron( speye( (p.N+1)^(p.nm-1) ), p.success_p) ;
                else
                    S =  kron( speye( (p.N+1)^(p.nm-1) ), p.success_n) ;
                end
            else
                if min_config(loop,j) == 1
                    S =  kron( speye( (p.N+1)^(j-1) ), p.success_p) ;
                else
                    S =  kron( speye( (p.N+1)^(j-1) ), p.success_n) ;
    
                end
                S =  kron( S, speye( (p.N+1)^(p.nm-j) )) ;
            end
            psi_update = S*psi_update;
        end
        P_sc = sum( conj(psi).*(psi_update),1 );
        suc_rate = suc_rate + P_sc;
    end
end

%% --- Purity Observable ---
function [output_purity]=purity(psi,p)
    % Computes the purity from trajectories
    density = sparse(psi*psi');
    density_op = density/size(psi,2);
    output_purity = trace(density_op*density_op);
end

%% --- Dissipator for Coupling ---
function out = coupling_generation(i,p)
    % Constructs dissipator L = a_i - J_ij a_j for mode coupling
    J_matrix = p.J(p);
    out = p.a{i(1)} - J_matrix(i(1),i(2))*p.a{i(2)} ; 
end

%% --- Initial State Generator ---
function psi_0 = initial(in)
    % cat state product
    alpha = in.alph; % value for alpha
    n = 0:in.N;
    coeff = (alpha.^n)./(factorial(n)).^.5;
    state_alpha = (exp(-.5*abs(alpha).^2).*coeff*eye(length(n),length(n))).';

    alpha_p = -in.alph;
    coeff_p = (alpha_p.^n)./(factorial(n)).^.5;
    state_alpha_p = (exp(-.5*abs(alpha_p).^2).*coeff_p*eye(length(n),length(n))).';

    state = state_alpha + state_alpha_p;
    total = kron(state,state);
    if in.nm>2
        for i = 3:in.nm
            total = kron(total,state);
        end
    end
    norm = sum(conj(total).*total);
    psi_0 = total./sqrt(norm);
end

%% --- Coupling Matrix Generator ---
function [J] = coupling_matrix(in)
    % Builds nearest-neighbor coupling matrix
    J = zeros(in.nm,in.nm);
        for i = 1:in.nm-1
           J(i,i+1) = 1;
           J(i+1,i) = J(i,i+1);
           if i == 1
              J(i,in.nm) = 1;
              J(in.nm,1) = J(i,in.nm);
           end
        end
end

%% --- Ising Configuration Generator ---
function [E_Ising,total_spin_config] = Ising_config(in,J)
    % Computes all spin configurations and their energies under Ising model
    spin_config = [1;-1];
    nm = in.nm;
    for l = 1:nm
        v = repelem(spin_config,2^nm/(2^l));
        total_spin_config(:,l) = repmat(v,2^nm/length(v),1);
    end
    E_Ising = 0;
    % JJ = in.J;
    JJ = in.J(in);
    % compute the corresponding Ising energy
    for k = 1:nm
        if k == nm
            E_Ising = E_Ising - JJ(k,1)*total_spin_config(:,k).*total_spin_config(:,1);
        else
            E_Ising = E_Ising - JJ(k,k+1)*total_spin_config(:,k).*total_spin_config(:,k+1);
        end
    end
    E_Ising = E_Ising*2; %% for symmetrical coupling
end

%% --- Success Projection Operators ---
function [success_p,success_n] = success_prob(nc)
    % Constructs projection matrices for detecting correct spin states
    % Uses Hermite polynomial integrals and Watson's formula for overlap
    success_n = zeros(nc+1,nc+1);
    m = 0:nc; %% ranges from 0 to the cutoff number. Indexed by k
    for k = 1:length(m)
        n = 0:1:m(k)-1; %% n goes from 0 to m - 1. Indexed by i
        for i = 1:length(n)
            coeff_n = 0; %% initialize the integrated integral value from - infinity to 0
            r = 0:n(i); %% r goes from 0 to n. Indexed by j. See paper by Watson, Journal of math soc, 1938
            for j = 1:length(r)
                var = m(k)-n(i)+2*r(j);
                if mod(var,2) == 0
                    integral_H_neg = 0; %% if var is even, the integral of H_{m-n+2r}*exp(-x^2), with limit 0 to infinity, is zero
                else              
                    numerator_neg = -2^(var-1)*sqrt(pi);
                    denominator_neg = gamma(1-.5*var);
                    integral_H_neg = numerator_neg/denominator_neg/( sqrt(pi)*sqrt( factorial(m(k)) * factorial(n(i)))*(sqrt(2))^(m(k)+n(i)) );
                end               
                coeff_temp_n = 2^n(i)*factorial(n(i))*nchoosek(m(k),n(i)-r(j))/(2^r(j)*factorial(r(j)))*integral_H_neg;
                coeff_n = coeff_n + coeff_temp_n;               

            end
            success_n(k,i) = coeff_n;
            success_n(i,k) = success_n(k,i);           
        end

    end
    % Build final sparse projectors
    success_p = -success_n + diag( .5*ones(1,nc+1) );
    success_n = success_n + diag( .5*ones(1,nc+1));
    success_p = sparse(success_p);
    success_n = sparse(success_n);
end
